#
import sys
from typing import Dict
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.figure import Figure
import torch
from torch.optim.lr_scheduler import ReduceLROnPlateau
from rsp.antenna.array import linear_array
#
from apps.wfs.wfs_config import WfsConfig as WG
from apps.wfs.wfs_mlp import WfsMlp
from apps.wfs.tcr_mlp import TcrMlp

class WfsApp(object):
    def __init__(self):
        self.name = 'apps.wfs.wfs_app.WfsApp'
    
    @staticmethod
    def startup(params: Dict = {}) -> None:
        print(f'波形优化程序 v0.0.5')
        WG.initialize()
        # WfsApp.t001()
        # WfsApp.waveform_demo(params=params)
        # WfsApp.test_dss(params=params)
        # 训练
        params['pt'] = None
        WfsApp.evaluate(params=params)
        WfsApp.train_main(params=params)
        # # 评估
        params['pt'] = './work/wfs/wfs.pt'
        WfsApp.evaluate(params=params)

    @staticmethod
    def train_main(params: Dict = {}) -> None:
        # 载入波束成形相关内容和训练目标
        X = torch.tensor([[WG.N, WG.f, WG.d, WG.theta]], dtype=torch.float64)
        # y_label = torch.load('./work/wfs/y_af.pt', weights_only=True)
        # for idx in range(y_label.shape[0]):
        #     if idx > 475 and idx < 535:
        #         y_label[idx] = 0.0
        #     else:
        #         y_label[idx] = -75.0
        # y_label = y_label.float()
        y_raw = torch.tensor([[0.26221649, 0.51874705, 0.81196007, 1., 1., 0.81196007, 0.51874705, 0.26221649]], dtype=torch.float64)
        y_label = y_raw # torch.log(y_raw)
        # 定义模型
        # model = WfsMlp()
        model = TcrMlp()
        # 转移到GPU
        X = X.to(WG.device)
        y_raw = y_raw.to(WG.device)
        y_label = y_label.to(WG.device)
        model = model.to(WG.device)
        WG.thetas = WG.thetas.to(WG.device)
        # 定义训练过程
        # loss_fn = torch.nn.MSELoss()
        loss_fn = torch.nn.L1Loss()
        # loss_fn = torch.nn.BCELoss()
        # 预热阶段
        print(f'##### 开始预热训练')
        warmup_opt = torch.optim.AdamW(model.parameters(), lr=1e-9, betas=(0.9, 0.99), weight_decay=1e-2)
        warmup_scheduler = ReduceLROnPlateau(warmup_opt, 'min', threshold=1e-4, eps=1e-8)
        warmup_epochs = 5
        for epoch in range(warmup_epochs):
            WfsApp.train_epoch(epoch, model, loss_fn, warmup_opt, warmup_scheduler, X, y_label, WG.N, WG.theta, WG.d, WG.f, WG.thetas)
        # 正式训练过程
        print(f'##### 开始正式训练')
        optimizer = torch.optim.AdamW(model.parameters(), lr=1e-7, betas=(0.9, 0.99), weight_decay=1e-2)
        scheduler = ReduceLROnPlateau(optimizer, 'min', threshold=1e-4, eps=1e-8)
        epochs = 200000
        best_loss = sys.float_info.max
        improve_threshold = 0.000001
        cumulative_steps = 0
        max_unimproves = 100
        pt_fn = './work/wfs/wfs.pt'
        for epoch in range(epochs):
            model.train()
            loss = WfsApp.train_epoch(epoch, model, loss_fn, optimizer, scheduler, X, y_label, WG.N, WG.theta, WG.d, WG.f, WG.thetas)
            print(f'epoch_{epoch}: loss={loss};')
            cumulative_steps += 1
            if best_loss > loss:
                if best_loss - loss > improve_threshold:
                    torch.save(model, pt_fn)
                    cumulative_steps = 0
                    best_loss = loss
            if cumulative_steps > max_unimproves:
                y = model(X)
                print(f'### y: {y}; ????????????????????????????????')
                print(f'### y: {y_raw}; ??????????????????????????????????????????')
                print(f'Earlly Stopping!!!!!!')
                break
        # 测试网络效果



    @staticmethod
    def train_epoch(epoch, model, loss_fn, optimizer, scheduler, X, y_label, N, theta, d, f, thetas):
        batch_size = 100
        total_loss = torch.tensor([0.0]).to(X.device)
        for bi in range(batch_size):
            # train
            y = model(X)
            # kwargs = {'number_of_elements': N,
            #         'scan_angle': theta /180.0 * np.pi,
            #         'element_spacing': d,
            #         'frequency': f,
            #         'theta': thetas,
            #         'coefficients': y[0]}
            # af = linear_array.dl_array_factor(**kwargs)
            # AF = 20.0 * torch.log10(torch.abs(af) + torch.finfo(float).eps)
            # loss = loss_fn(y_label, AF)

            loss = loss_fn(y, y_label)

            # loss += torch.sum(torch.max(cmp_base, low_limit - y) + torch.max(cmp_base, y - high_limit))
            total_loss += loss
            # loss.backward()
            # optimizer.step()
            # optimizer.zero_grad()
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
        if scheduler is not None:
            scheduler.step(total_loss)
        return total_loss.detach().cpu().item()/batch_size

        

    @staticmethod
    def evaluate(params: Dict = {}) -> None:
        # 定义模型
        if params['pt'] is None:
            # model = WfsMlp()
            model = TcrMlp()
        else:
            model = torch.load(params['pt'])
        model.eval()
        model = model.to(WG.device)
        X = torch.tensor([[WG.N, WG.f, WG.d, WG.theta]], dtype=torch.float64)
        X = X.to(WG.device)
        y = model(X)
        kwargs = {'number_of_elements': WG.N,
                'scan_angle': WG.theta /180.0 * np.pi,
                'element_spacing': WG.d,
                'frequency': WG.f,
                'theta': WG.thetas,
                'coefficients': y[0]}
        af_pt = linear_array.dl_array_factor(**kwargs)
        af = af_pt.detach().cpu().numpy()
        AF_pt = 20.0 * torch.log10(torch.abs(af_pt) + torch.finfo(float).eps)
        AF = AF_pt.detach().cpu().numpy()
        # print(f'y_af: {y_af.shape}; {y_af.dtype};')
        # af = y_af.numpy()
        bw, sll = linear_array.calculate_bw_sll(20.0 * np.log10(abs(af) + np.finfo(float).eps), WG.x_scale)
        print(f'波束宽度：{bw}; 旁瓣水平：{sll}; N: {WG.N}; Theta: {WG.theta};')
        fig = plt.figure(1) #Figure()
        axes1 = fig.add_subplot(111)
        axes1.clear()
        thetas_np = WG.thetas.detach().cpu().numpy()
        axes1.plot(np.degrees(thetas_np), 20.0 * np.log10(abs(af) + np.finfo(float).eps), '')
        axes1.set_ylim(-80, 5)
        axes1.set_xlabel("Theta (degrees)", size=12)
        axes1.set_ylabel("Array Factor (dB)", size=12)
        axes1.grid(linestyle=':', linewidth=0.5)
        axes1.set_title('Linear Array Antenna Pattern', size=14)
        axes1.tick_params(labelsize=12)
        plt.show()
        print(f'^_^ The End! ^_^')




        

    @staticmethod
    def t001():
        y_label = torch.load('./work/wfs/y_af.pt')
        for idx in range(y_label.shape[0]):
            if idx > 475 and idx < 535:
                y_label[idx] = 0.0
            else:
                y_label[idx] = -45.0
        y_label = y_label.detach().cpu().numpy()
        print(f'y_label: {y_label.shape};')
        fig = plt.figure(1) #Figure()
        axes1 = fig.add_subplot(111)
        axes1.clear()
        thetas_np = WG.thetas.detach().cpu().numpy()
        axes1.plot(np.degrees(thetas_np), y_label, '')
        axes1.set_ylim(-80, 5)
        axes1.set_xlabel("Theta (degrees)", size=12)
        axes1.set_ylabel("Array Factor (dB)", size=12)
        axes1.grid(linestyle=':', linewidth=0.5)
        axes1.set_title('Linear Array Antenna Pattern', size=14)
        axes1.tick_params(labelsize=12)
        plt.show()





    @staticmethod
    def test_dss(params:Dict = {}) -> None:
        print(f'测试载入数据集')
        N = 32 # 线性天线个数
        f = 300.0E6 # 300MHz频率
        c = 3E8 # 光速
        lambda_ = c / f
        d = 0.5*lambda_ # 天线间距离
        Theta = 5 / 180.0 * torch.pi # 扫描角度，单位为弧度
        x_min, x_max, x_nums = -torch.pi/2.0, torch.pi/2.0, 1000
        thetas = np.linspace(x_min, x_max, x_nums) # 定义绘制扫描角度与天线因子图时横坐标扫描角度的范围
        x_scale = (x_max - x_min) / x_nums
        # 载入AF
        y_af = torch.load('./work/wfs/y_af.pt')
        # y_af = y_af.float()
        print(f'y_af: {y_af.shape}; {y_af.dtype};')
        af = y_af.numpy()
        bw, sll = linear_array.calculate_bw_sll(20.0 * np.log10(abs(af) + np.finfo(float).eps), x_scale)
        print(f'波束宽度：{bw}; 旁瓣水平：{sll}; N: {N}; Theta: {Theta};')
        fig = plt.figure(1) #Figure()
        axes1 = fig.add_subplot(111)
        axes1.clear()
        axes1.plot(np.degrees(thetas), 20.0 * np.log10(abs(af) + np.finfo(float).eps), '')
        axes1.set_ylim(-80, 5)
        axes1.set_xlabel("Theta (degrees)", size=12)
        axes1.set_ylabel("Array Factor (dB)", size=12)
        axes1.grid(linestyle=':', linewidth=0.5)
        axes1.set_title('Linear Array Antenna Pattern', size=14)
        axes1.tick_params(labelsize=12)
        plt.show()
        print(f'^_^ The End! ^_^')

    @staticmethod
    def waveform_demo(params: Dict = {}) -> None:
        N = 32 # 线性天线个数
        f = 300.0E6 # 300MHz频率
        c = 3E8 # 光速
        lambda_ = c / f
        print(f'波长为：lambdam={lambda_};')
        d = 0.5*lambda_ # 天线间距离
        Theta = 5 / 180.0 * torch.pi # 扫描角度，单位为弧度
        x_min, x_max, x_nums = -torch.pi/2.0, torch.pi/2.0, 1000
        thetas = np.linspace(x_min, x_max, x_nums) # 定义绘制扫描角度与天线因子图时横坐标扫描角度的范围
        x_scale = (x_max - x_min) / x_nums
        # 计算天线因子
        # Set up the args
        antenna_type = 'Tschebyscheff'
        side_lobe_level = -30.0
        kwargs = {'number_of_elements': N,
                  'scan_angle': Theta,
                  'element_spacing': d,
                  'frequency': f,
                  'theta': thetas,
                  'window_type': antenna_type,
                  'side_lobe_level': side_lobe_level}
        af = linear_array.array_factor(**kwargs)
        print(f'raw AF: {af.shape}; {af.dtype};')
        y_af = torch.from_numpy(af)
        torch.save(y_af, './work/wfs/y_af.pt')
        bw, sll = linear_array.calculate_bw_sll(20.0 * np.log10(abs(af) + np.finfo(float).eps), x_scale)
        print(f'波束宽度：{bw}; 旁瓣水平：{sll}; N: {N}; Theta: {Theta};')
        fig = plt.figure(1) #Figure()
        axes1 = fig.add_subplot(111)
        axes1.clear()
        axes1.plot(np.degrees(thetas), 20.0 * np.log10(abs(af) + np.finfo(float).eps), '')
        axes1.set_ylim(-80, 5)
        axes1.set_xlabel("Theta (degrees)", size=12)
        axes1.set_ylabel("Array Factor (dB)", size=12)
        axes1.grid(linestyle=':', linewidth=0.5)
        axes1.set_title('Linear Array Antenna Pattern', size=14)
        axes1.tick_params(labelsize=12)
        plt.show()
        print(f'^_^ The End! ^_^')